import ctypes
import os
import time

# --- Configuration Section ---
# Update DLL name and path according to your setup
DLL_NAME = "ssh.dll"
DLL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), DLL_NAME)

# Your SSH Server Information
SSH_HOST = "192.168.31.233"
SSH_PORT = 22
SSH_USER = "zyu"
SSH_PASS = "6"

# --- Load DLL and Define Function Interfaces ---
try:
    ssh_dll = ctypes.cdll.LoadLibrary(DLL_PATH)
except Exception as e:
    print(f"Error loading DLL: {e}")
    exit(1)

# Define the interface for the Connect function
# int Connect(char* host, int port, char* user, char* password)
ssh_dll.Connect.argtypes = [ctypes.c_char_p, ctypes.c_int, ctypes.c_char_p, ctypes.c_char_p]
ssh_dll.Connect.restype = ctypes.c_int

# Define the interface for the SendCommand function
# char* SendCommand(char* command)
ssh_dll.SendCommand.argtypes = [ctypes.c_char_p]
ssh_dll.SendCommand.restype = ctypes.c_char_p

# Define the interface for the Disconnect function
# void Disconnect()
ssh_dll.Disconnect.argtypes = []
ssh_dll.Disconnect.restype = None

def test_ssh_dll():
    """
    Connects once, sends multiple commands, and then disconnects.
    """
    print("--- 1. Attempting to connect to SSH server ---")
    host_c = SSH_HOST.encode('utf-8')
    user_c = SSH_USER.encode('utf-8')
    pass_c = SSH_PASS.encode('utf-8')

    result = ssh_dll.Connect(host_c, SSH_PORT, user_c, pass_c)

    if result == 0:
        print(f"✅ Connection successful to {SSH_HOST}:{SSH_PORT}")
        time.sleep(1) # Wait for one second to ensure a stable connection

        # List of commands to execute
        commands = [
            "ls -l",
            "pwd",
            "whoami",
            "uname -a"
        ]

        for cmd in commands:
            print(f"\n--- Sending command: {cmd} ---")
            command_c = cmd.encode('utf-8')
            output_ptr = ssh_dll.SendCommand(command_c)

            if output_ptr:
                try:
                    output_str = ctypes.cast(output_ptr, ctypes.c_char_p).value.decode('utf-8')
                    print("✅ Command executed successfully. Output:")
                    print("---------------------------------------")
                    print(output_str)
                    print("---------------------------------------")
                except UnicodeDecodeError:
                    print("❌ Failed to decode output. Non-UTF-8 characters might be present.")
                finally:
                    # Free the memory allocated by the Go code
                    ctypes.CDLL('msvcrt').free(output_ptr)
            else:
                print("❌ Command execution failed, no data returned.")
                # We do not break here, in case the next command succeeds.
    else:
        print("❌ Connection failed. Please check IP, port, username, and password.")

    print("\n--- 3. Disconnecting ---")
    ssh_dll.Disconnect()
    print("✅ Disconnection successful.")


# --- Run the test ---
if __name__ == "__main__":
    test_ssh_dll()
